package com.piece.core.web.authority.token;

import com.piece.core.framework.constant.ExceptionAuthConstants;
import com.piece.core.framework.constant.SecurityConstants;
import com.piece.core.framework.util.basic.I18nUtil;
import com.piece.core.framework.util.math.SequenceUtil;
import com.piece.core.framework.util.string.StringUtil;
import com.piece.core.web.properties.SecurityProperties;
import com.piece.core.framework.support.dto.UserDetailDTO;
import org.apache.commons.collections.MapUtils;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.oauth2.common.*;
import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.exceptions.InvalidScopeException;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.common.exceptions.UnapprovedClientAuthenticationException;
import org.springframework.security.oauth2.provider.*;
import org.springframework.security.oauth2.provider.token.DefaultTokenServices;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.transaction.annotation.Transactional;
import java.util.Collection;
import java.util.Date;
import java.util.Set;

/**
 * 重写 DefaultTokenServices，实现登录同应用同账号互踢
 */
public class TokenServices extends DefaultTokenServices {

    private AuthenticationManager authenticationManager;
    private TokenStore tokenStore;
    private TokenEnhancer accessTokenEnhancer;
    private ClientDetailsService clientDetailsService;
    private boolean supportRefreshToken = false;
    private boolean reuseRefreshToken = true;
    private boolean kickLogin;
    private int maximum;

    public TokenServices(boolean kickLogin, SecurityProperties securityProperties) {
        this.kickLogin = kickLogin;
        this.maximum = securityProperties.getToken().getMaximum();
    }

    @Override
    @Transactional
    public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication) throws AuthenticationException {
        OAuth2AccessToken existingAccessToken = tokenStore.getAccessToken(authentication);
        OAuth2AccessToken kickAccessToken;
        OAuth2RefreshToken refreshToken = null;

        if (null != existingAccessToken) {
            if (existingAccessToken.isExpired()) {
                if (null != existingAccessToken.getRefreshToken()) {
                    refreshToken = existingAccessToken.getRefreshToken();
                    tokenStore.removeRefreshToken(refreshToken);
                }
            } else {
                tokenStore.storeAccessToken(existingAccessToken, authentication);
                return existingAccessToken;
            }
        }

        if (kickLogin) {
            Collection<OAuth2AccessToken> tokens =
                    tokenStore.findTokensByClientIdAndUserName(authentication.getOAuth2Request().getClientId(), authentication.getName());
            if (null != tokens && tokens.size() >= maximum) {
                kickAccessToken = tokens.iterator().next();
                if (null != kickAccessToken.getRefreshToken()) {
                    tokenStore.removeRefreshToken(kickAccessToken.getRefreshToken());
                }
                tokenStore.removeAccessToken(kickAccessToken);
            }
        }

        if (null == refreshToken) {
            refreshToken = createRefreshToken(authentication);
        } else if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
            ExpiringOAuth2RefreshToken expiring = (ExpiringOAuth2RefreshToken) refreshToken;
            if (System.currentTimeMillis() > expiring.getExpiration().getTime()) {
                refreshToken = createRefreshToken(authentication);
            }
        }

        OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
        tokenStore.storeAccessToken(accessToken, authentication);

        refreshToken = accessToken.getRefreshToken();
        if (null != refreshToken) {
            tokenStore.storeRefreshToken(refreshToken, authentication);
        }
        return accessToken;
    }

    private OAuth2RefreshToken createRefreshToken(OAuth2Authentication authentication) {
        if (!isSupportRefreshToken(authentication.getOAuth2Request())) {
            return null;
        }

        int validitySeconds = getRefreshTokenValiditySeconds(authentication.getOAuth2Request());
        String value = SequenceUtil.uuid();
        if (validitySeconds > 0) {
            return new DefaultExpiringOAuth2RefreshToken(value, new Date(System.currentTimeMillis()
                    + (validitySeconds * 1000L)));
        }
        return new DefaultOAuth2RefreshToken(value);
    }

    private OAuth2AccessToken createAccessToken(OAuth2Authentication authentication, OAuth2RefreshToken refreshToken) {
        DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(SequenceUtil.uuid());
        int validitySeconds = getAccessTokenValiditySeconds(authentication.getOAuth2Request());
        if (validitySeconds > 0) {
            token.setExpiration(new Date(System.currentTimeMillis() + (validitySeconds * 1000L)));
        }
        token.setRefreshToken(refreshToken);
        token.setScope(authentication.getOAuth2Request().getScope());

        return (null != accessTokenEnhancer) ? accessTokenEnhancer.enhance(token, authentication) : token;
    }

    @Override
    @Transactional(noRollbackFor = {InvalidTokenException.class, InvalidGrantException.class})
    public OAuth2AccessToken refreshAccessToken(String refreshTokenValue, TokenRequest tokenRequest) throws AuthenticationException {
        if (!supportRefreshToken) {
            throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);
        }

        OAuth2RefreshToken refreshToken = tokenStore.readRefreshToken(refreshTokenValue);
        if (null == refreshToken) {
            throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);
        }

        OAuth2Authentication authentication = tokenStore.readAuthenticationForRefreshToken(refreshToken);
        if (null != authenticationManager && !authentication.isClientOnly()) {
            authenticationManager.authenticate(authentication.getUserAuthentication());
        }

        String clientId = authentication.getOAuth2Request().getClientId();
        if (null == clientId || !clientId.equals(tokenRequest.getClientId())) {
            throw new InvalidGrantException("Wrong client for this refresh token: " + refreshTokenValue);
        }
        tokenStore.removeAccessTokenUsingRefreshToken(refreshToken);

        if (isExpired(refreshToken)) {
            tokenStore.removeRefreshToken(refreshToken);
            throw new InvalidTokenException("Expired refresh token: " + refreshToken);
        }

        authentication = createRefreshedAuthentication(authentication, tokenRequest);
        if (!reuseRefreshToken) {
            tokenStore.removeRefreshToken(refreshToken);
            refreshToken = createRefreshToken(authentication);
        }

        OAuth2AccessToken accessToken = createAccessToken(authentication, refreshToken);
        tokenStore.storeAccessToken(accessToken, authentication);
        if (!reuseRefreshToken) {
            tokenStore.storeRefreshToken(accessToken.getRefreshToken(), authentication);
        }
        return accessToken;
    }

    private OAuth2Authentication createRefreshedAuthentication(OAuth2Authentication authentication, TokenRequest request) {
        OAuth2Authentication narrowed;
        Set<String> scope = request.getScope();
        OAuth2Request clientAuth = authentication.getOAuth2Request().refresh(request);
        if (null != scope && !scope.isEmpty()) {
            Set<String> originalScope = clientAuth.getScope();
            if (null == originalScope || !originalScope.containsAll(scope)) {
                throw new InvalidScopeException("Unable to narrow the scope of the client authentication to " + scope
                        + ".", originalScope);
            } else {
                clientAuth = clientAuth.narrowScope(scope);
            }
        }
        narrowed = new OAuth2Authentication(clientAuth, authentication.getUserAuthentication());
        return narrowed;
    }

    public OAuth2AccessToken getUserTokenInfo(String clientId, UserDetailDTO userDetail) {
        OAuth2AccessToken oauth2AccessToken;
        ClientDetails clientDetails = clientDetailsService.loadClientByClientId(clientId);
        if (null == clientDetails) {
            throw new UnapprovedClientAuthenticationException(I18nUtil.message(ExceptionAuthConstants.CLIENT_INVALID));
        }

        TokenRequest tokenRequest = new TokenRequest(MapUtils.EMPTY_MAP, clientId, clientDetails.getScope(), SecurityConstants.USERNAME_CODE);
        OAuth2Request oAuth2Request = tokenRequest.createOAuth2Request(clientDetails);
        Authentication authToken;
        if (StringUtil.isNotEmpty(userDetail.getOpenID())) {
            authToken = new OpenIdAuthenticationToken(userDetail.getOpenID(), clientId);
        } else {
            authToken = new UsernameAuthenticationToken(userDetail.getLoginName(),userDetail.getPassword(), clientId);
        }

        Authentication authentication = authenticationManager.authenticate(authToken);
        OAuth2Authentication oauth2Authentication = new OAuth2Authentication(oAuth2Request, authentication);
        oauth2AccessToken = createAccessToken(oauth2Authentication);
        oauth2Authentication.setAuthenticated(true);
        SecurityContextHolder.getContext().setAuthentication(oauth2Authentication);

        return oauth2AccessToken;
    }

    @Override
    public void setAuthenticationManager(AuthenticationManager authenticationManager) {
        this.authenticationManager = authenticationManager;
        super.setAuthenticationManager(authenticationManager);
    }

    @Override
    public void setTokenStore(TokenStore tokenStore) {
        this.tokenStore = tokenStore;
        super.setTokenStore(tokenStore);
    }

    @Override
    public void setTokenEnhancer(TokenEnhancer accessTokenEnhancer) {
        this.accessTokenEnhancer = accessTokenEnhancer;
        super.setTokenEnhancer(accessTokenEnhancer);
    }

    @Override
    public void setClientDetailsService(ClientDetailsService clientDetailsService) {
        super.setClientDetailsService(clientDetailsService);
        this.clientDetailsService = clientDetailsService;
    }

    @Override
    public void setSupportRefreshToken(boolean supportRefreshToken) {
        this.supportRefreshToken = supportRefreshToken;
        super.setSupportRefreshToken(supportRefreshToken);
    }

    @Override
    public void setReuseRefreshToken(boolean reuseRefreshToken) {
        this.reuseRefreshToken = reuseRefreshToken;
        super.setReuseRefreshToken(reuseRefreshToken);
    }

    public void setKickLogin(boolean kickLogin) {
        this.kickLogin = kickLogin;
    }

    public void setMaximum(int maximum) {
        this.maximum = maximum;
    }
}
